#include <cstdio>
#include <vector>
#include <cuda_runtime.h>
#include <curand_kernel.h>

__device__ bool Iterate(float2 c0, int *n, float2 *seq) {
  float2 c = make_float2(0.0f, 0.0f), cn;
  *n = 0;
  for (int i = 0; i < 200; ++i) {
    cn.x = c.x * c.x - c.y * c.y + c0.x;
    cn.y = 2.0f * c.x * c.y + c0.y;
    seq[i].x = cn.x;
    seq[i].y = cn.y;
    if (cn.x*cn.x + cn.y*cn.y > 10.0f) {
      *n = i;
      return true;
    }
    c.x = cn.x;
    c.y = cn.y;
  }
  return false;
}
__global__ void Init(curandState *state) {
  const int idx = blockIdx.x * blockDim.x + threadIdx.x;
  curand_init(1234, idx, 0, &state[idx]);
}
__global__ void Buddhabrot(int size, unsigned int *image, curandState *state) {
  const int idx = blockIdx.x * blockDim.x + threadIdx.x;
  float2 c;
  float2 seq[200];
  int n;
  curandState local_state = state[idx];
  for (int i = 0; i < 1000; ++i) {
    c.x = curand_uniform(&local_state);
    c.y = curand_uniform(&local_state);
    c.x = 4.2f * c.x - 2.1f;
    c.y = 4.2f * c.y - 2.1f;
    if (Iterate(c, &n, seq)) {
      for (int j = 0; j < n; ++j) {
        int ix = 0.3f * size * (seq[j].x + 0.5f) + size/2;
        int iy = 0.3f * size * seq[j].y + size/2;
        if (ix >= 0 && iy >= 0 && ix < size && iy < size) {
          atomicAdd(&image[iy * size + ix], 1);
        }
      }
    }
  }
  state[idx] = local_state;
}

int main(int argc, char *argv[]) {
  int size = 800;

  curandState *dev_states;
  cudaMalloc(&dev_states, 512 * 256 * sizeof(curandState));
  Init<<<512, 256>>>(dev_states);

  unsigned int *image_dev;
  cudaMalloc(&image_dev, size * size * sizeof(unsigned int));
  cudaMemset(image_dev, 0, size * size * sizeof(unsigned int));

  for (int i = 0; i < 8; ++i) {
    Buddhabrot<<<512, 200>>>(size, image_dev, dev_states);
    cudaDeviceSynchronize();
  }
  
  std::vector<unsigned int> image_host(size * size);
  cudaMemcpy(&image_host[0], image_dev, size * size * sizeof(unsigned int),
             cudaMemcpyDeviceToHost);

  fprintf(stderr, "error: %s\n", cudaGetErrorString(cudaGetLastError()));

  unsigned int minval, maxval = 0;
  for (int i = 0; i < size*size; ++i) {
    maxval = std::max(maxval, image_host[i]);
  }
  minval = maxval;
  for (int i = 0; i< size*size; ++i) {
    minval = std::min(minval, image_host[i]);
  }
  fprintf(stderr, "Density value range: %d to %d\n", minval, maxval);

  FILE *fp = fopen("buddhabrot.tga", "wb");
  const uchar1 tga_header[18] = {
    0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0,
    size & 0x00ff, (size & 0xff00) >> 8,
    size & 0x00ff, (size & 0xff00) >> 8,
    32, 0
  };
  fwrite(tga_header, sizeof(tga_header), 1, fp);
  // Raw uncompressed bytes
  const float r = static_cast<float>(maxval) - minval;
  for (int i = 0; i < size*size; ++i) {
    float ramp = 2*(image_host[i] - minval) / r;
    if (ramp > 1)
      ramp = 1;
    ramp = pow(ramp,0.5);
    fputc(255,fp);
    fputc((int)(ramp*255),fp);
    fputc((int)(ramp*255),fp);
    fputc(255,fp);
  }
  fclose(fp);

  return 0;
}
